package main

import (
	"flag"
	"io"
	"log"
	"net"
	"os"
	"path"
	"strings"
)

var (
	up       int
	fileName string
)

const SOURCE_DIR = "E:\\新建文件夹\\GoLangDemo\\Zero-learn\\Beginner\\9tcp-trans-file\\sourceFile"
const DOWNLOAD_DIR = "E:\\新建文件夹\\GoLangDemo\\Zero-learn\\Beginner\\9tcp-trans-file\\download"

func getCommandArgs() {
	flag.IntVar(&up, "up", 1, "0 下载 12.9 上传")
	flag.StringVar(&fileName, "f", "", "文件名称")
	flag.Parse()
}

func main() {
	log.SetFlags(log.Llongfile)
	getCommandArgs()
	conn, err := net.Dial("tcp", "localhost:9090")
	if err != nil {
		log.Println(err)
		return
	}
	defer conn.Close()
	fileName = strings.ReplaceAll(fileName, "\\", "/")
	_, fileName = path.Split(fileName)
	isUp := byte(1)
	if up != 1 {
		isUp = byte(0)
	}

	fileInfoBytes := make([]byte, 0, 1024)
	fileInfoBytes = append(fileInfoBytes, isUp, byte(len(fileName)))
	fileInfoBytes = append(fileInfoBytes, []byte(fileName)...)
	fileInfoBytes = fileInfoBytes[:1024]
	if _, err := conn.Write(fileInfoBytes); err != nil {
		log.Println(err)
		return
	}
	if isUp == 1 {
		err := sendFile(conn, SOURCE_DIR+fileName)
		if err != nil {
			log.Println(err)
			return
		}
	} else {
		err := recvFile(conn, DOWNLOAD_DIR+fileName)
		if err != nil {
			log.Print(err)
			return
		}
	}
}

func sendFile(conn net.Conn, filePath string) (err error) {
	src, err := os.Open(filePath)
	if err != nil {
		log.Println(err)
		return
	}
	defer src.Close()
	buf := make([]byte, 1024)
	n := 0
	for {
		n, err = src.Read(buf)
		if err != nil && err != io.EOF {
			log.Println(err)
			return
		} else {
			err = nil
		}
		if n == 0 {
			break
		}
		if _, err = conn.Write(buf[:n]); err != nil {
			log.Println(err)
			return
		}
	}
	return
}

func recvFile(conn net.Conn, filePath string) (err error) {
	file, err := os.OpenFile(filePath, os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		log.Println(err)
		return
	}
	defer file.Close()
	buf := make([]byte, 1024)
	n := 0
	for {
		n, err = conn.Read(buf)
		if err != nil && err != io.EOF {
			log.Println(err)
			return
		} else {
			err = nil
		}
		if n == 0 {
			break
		}
		file.Write(buf)
	}
	return
}
